/**
 * 
 */
package com.gr.experimental.model.impl.distr;

import java.util.ArrayList;
import java.util.Arrays;
import java.util.Collection;
import java.util.HashMap;
import java.util.Iterator;
import java.util.List;
import java.util.Map;

import org.antlr.runtime.ANTLRStringStream;
import org.antlr.runtime.CommonTokenStream;
import org.antlr.runtime.RecognitionException;
import org.antlr.runtime.tree.Tree;

import com.gr.experimental.antlr.LDDLexer;
import com.gr.experimental.antlr.LDDParser;
import com.gr.experimental.model.Distribution;
import com.gr.experimental.model.MEBNConstants;
import com.gr.experimental.model.MTheoryManager;
import com.gr.experimental.model.RVClass;
import com.gr.experimental.model.RVState;
import com.gr.experimental.model.UID;
import com.gr.experimental.model.impl.distr.operations.AssignAction;
import com.gr.experimental.model.impl.distr.operations.BinaryOperation;
import com.gr.experimental.model.impl.distr.operations.Expression;
import com.gr.experimental.model.impl.distr.operations.Function;
import com.gr.experimental.model.impl.distr.operations.FunctionCall;
import com.gr.experimental.model.impl.distr.operations.LiteralExpression;
import com.gr.experimental.model.impl.distr.operations.ReturnAction;
import com.gr.experimental.model.impl.distr.operators.BinaryOperator;
import com.gr.experimental.model.impl.distr.operators.DivOperator;
import com.gr.experimental.model.impl.distr.operators.EqualityOperator;
import com.gr.experimental.model.impl.distr.operators.GreaterEqualOperator;
import com.gr.experimental.model.impl.distr.operators.GreaterThanOperator;
import com.gr.experimental.model.impl.distr.operators.LogicalAndOperator;
import com.gr.experimental.model.impl.distr.operators.LogicalOrOperator;
import com.gr.experimental.model.impl.distr.operators.LowerEqualOperator;
import com.gr.experimental.model.impl.distr.operators.LowerThanOperator;
import com.gr.experimental.model.impl.distr.operators.MinusOperator;
import com.gr.experimental.model.impl.distr.operators.PlusOperator;
import com.gr.experimental.model.impl.distr.operators.RestOperator;
import com.gr.experimental.model.impl.distr.operators.TimesOperator;

/**
 * @author Karol Banczyk
 *
 */
public class LDDBuilder {

	private Map<String, RVClass> rvClassesByNameMap;

	private Map<String, UID> allUidsByNameMap;

	private Map<String, RVClass> parentsByNameMap = new HashMap<String, RVClass>();
	
	private List<RVClass> parents = new ArrayList<RVClass>();
	
	String distributionId = null;
	
	public static final String RETURN_ARRAY_NAME = "prob";

	Action definition;
	List<UID> domain = new ArrayList<UID>();
	
	Map<String, Descriptor> allDescriptorMap = new HashMap<String, Descriptor>();
	Map<String, Descriptor> superContextDescriptorMap = new HashMap<String, Descriptor>();
	Descriptor returnDescriptor;
	
	void addDescriptor(Tree tree, Descriptor descriptor, boolean isReturnDescriptor) throws LDDCompilationException {
		if (superContextDescriptorMap.containsKey(descriptor.getSymbolName())) {
			throw new LDDCompilationException("Symbol " + descriptor.getSymbolName() 
					+ " redefined " + createLineInformation(tree));
		}
		if (isReturnDescriptor) {
			if (returnDescriptor != null) {
				throw new IllegalArgumentException("Return descriptor already set");
			}
			returnDescriptor = descriptor;
		} else {
			superContextDescriptorMap.put(descriptor.getSymbolName(), descriptor);
		}
		allDescriptorMap.put(descriptor.getSymbolName(), descriptor);
	}
	
	
	
	// main actions
//	List<Action> mainActions = new ArrayList<Action>();
	
	private Map<String, UID> domainUidsByName = new HashMap<String, UID>();
	
	LDDBuilder(MTheoryManager manager, List<RVClass> otherNodes, Collection<UID> uids) {
		this.rvClassesByNameMap = createRVClassByNameMap(otherNodes);
		List<UID> booleanUIDs = manager.getBuiltInElements().getBooleanUIDs();
		
		List<UID> allUIDs = new ArrayList<UID>();
		allUIDs.addAll(booleanUIDs);
		allUIDs.addAll(uids);
		
		this.allUidsByNameMap = createUidsByNameMap(allUIDs);
	}
	
	
	Map<String, RVClass> createRVClassByNameMap(List<RVClass> rvClasses) {
		Map<String, RVClass> retMap = new HashMap<String, RVClass>();
		for (Iterator<RVClass> it=rvClasses.iterator(); it.hasNext();){
			RVClass rvClass = it.next();
			retMap.put(rvClass.getName(), rvClass);
		}
		return retMap;
	}

	Map<String, UID> createUidsByNameMap(List<UID> uids) {
		Map<String, UID> retMap = new HashMap<String, UID>();
		for (Iterator<UID> it=uids.iterator(); it.hasNext();){
			UID uid = it.next();
			retMap.put(uid.getValue(), uid);
		}
		return retMap;
	}

	Distribution buildDistribution(String def) throws LDDCompilationException, LDDParseException {
		
		ANTLRStringStream inputStream = new ANTLRStringStream(def);
		
		LDDLexer lexer = new LDDLexer(inputStream);
		CommonTokenStream tokens = new CommonTokenStream(lexer);
		LDDParser parser = new LDDParser(tokens);
		
		Tree tree;
		try {
			tree = (Tree) parser.distribution().getTree();
		} catch (RecognitionException e) {
			throw new LDDParseException(
					"Parse error in line: " + e.line + 
					", col: " + e.charPositionInLine + ".", e
					);
		}
		
		// handle the AST tree and 
		// extract needed structures to instance variables
		// of this builder class
		handleDistributionSignature(tree);
		
		// create the distribution using 
		// the extracted structure
		//new ArrayList<RVClass>(this.parentsByNameMap.values());
		
		// the super context contains all unmodifialbe data (shared between 
		// function executions
		// it does not therefore contain the return descriptor
		// which exists only within a single execution
		ContextImpl superContext = new ContextImpl(superContextDescriptorMap, null);
		LocalDynamicDistribution distribution = 
			new LocalDynamicDistribution(
					parents, domain, definition, superContext, null, returnDescriptor
			);
		return distribution;
	}
		
	void handleDistributionSignature(Tree tree) throws LDDCompilationException {

		switch(tree.getType()) {
		case LDDParser.DISTRIBUTION:
			for (int i=0; i < tree.getChildCount(); i++) {
				handleDistributionSignature(tree.getChild(i));
			}
			break;
		case LDDParser.PARENTS:
			handleParents(tree);
			break;
		case LDDParser.DISTRIBUTION_ID:
			handleDistributionId(tree);
			break;
		case LDDParser.DOMAIN:
			handleDomain(tree);
			break;
		case LDDParser.DEFINITION:
			definition = handleAction(tree.getChild(0));
			break;
		default:
			throw createUnexpectedException(tree);
		}
	}
	
	Action handleAction(Tree tree) throws LDDCompilationException {
		switch(tree.getType()) {
		case LDDParser.IF:
			return handleConditionalAction(tree);
		case LDDParser.BLOCK:
			return handleBlock(tree);
		case LDDParser.ASSIGNMENT:
			return handleAssignment(tree);
		case LDDParser.RETURN:
			return handleReturn(tree);
		default:
			throw createUnexpectedException(tree);
		}
	}
	
	public Action handleReturn(Tree tree) throws LDDCompilationException {
		// we currently have 2 cases
		// either an empty return (with no parameters)
		// or a return with literal distribution specification

		
		if (tree.getChildCount() == 0) { // the first case
			return new ReturnAction();
		} else {
			// return a block with return value initialization
			// and return action as its final action
			
			Tree arrayValueTree = tree.getChild(0);
			assert arrayValueTree.getType() == LDDParser.ARRAY_VALUE;
			
			// we actually know that there may be only one return type
			// and there's currently no real need to make this place more general
			
			if (arrayValueTree.getChildCount()!= domain.size()) {
				throw new LDDCompilationException("Array size is " + 
						arrayValueTree.getChildCount() + 
						" and should be equal to domain size: " + domain.size() 
						+ " "  + createLineInformation(arrayValueTree));
			}
			
			List<Action> actions = new ArrayList<Action>();
			for (int i=0; i < arrayValueTree.getChildCount(); i++) {
				// if we can not cast to double
				Expression assignaedValueExpression = handleExpression(arrayValueTree.getChild(i));
				if (!Number.class.isAssignableFrom(assignaedValueExpression.getValueType())) {
					throw new LDDCompilationException(
							"Bad argument type " + createLineInformation(arrayValueTree.getChild(i)));
				}
				assert returnDescriptor != null;
				Assignable arrayAccess = new ArrayAccess(returnDescriptor, new LiteralExpression(i));
				actions.add(new AssignAction(arrayAccess, assignaedValueExpression));
			}
			actions.add(new ReturnAction());

			return new Block(actions);
		}

	}


	Action handleAssignment(Tree tree) throws LDDCompilationException {
		
		Assignable assignable = handleAssignable(tree.getChild(0)); // the assignable
		Expression assignedExpression = handleExpression(tree.getChild(1)); // the assigned expression's value

		return new AssignAction(assignable, assignedExpression) ;
	}
	
	Assignable handleAssignable(Tree tree) throws LDDCompilationException {
		switch(tree.getType()) {
		case LDDParser.ARRAY_ACCESS:
			return handleArrayAccess(tree);
		case LDDParser.IDENTIFIER:
			return new ContextAccess(handleIdentifier(tree));
		default:
			throw createUnexpectedException(tree);
		}
	}
	
	Expression handleExpression(Tree tree) throws LDDCompilationException {
		switch(tree.getType()) {
		case LDDParser.AND:
		case LDDParser.OR:
		case LDDParser.PLUS:
		case LDDParser.MINUS:
		case LDDParser.TIMES: 
		case LDDParser.DIV:
		case LDDParser.REST: 
		case LDDParser.LT:	
		case LDDParser.LTE:	
		case LDDParser.GT:	
		case LDDParser.GTE:	
		case LDDParser.EQUALS:
			return handleBinaryOperation(tree);
		case LDDParser.IDENTIFIER:
            return new ContextAccess(handleIdentifier(tree));
		case LDDParser.FUNCTION_CALL:
			return handleFunctionCall(tree);
		case LDDParser.ARRAY_ACCESS:
			return handleArrayAccess(tree);
		case LDDParser.DecimalLiteral:
		case LDDParser.FloatingPointLiteral:
		case LDDParser.TRUE:
		case LDDParser.FALSE:
		case LDDParser.ABSURD:
			return handleLiteral(tree);
		default:
			throw createUnexpectedException(tree);
		}
	}
	
	ArrayAccess handleArrayAccess(Tree tree) throws LDDCompilationException {
		Descriptor descriptor = handleIdentifier(tree.getChild(0));
		Expression indexExpression = handleExpression(tree.getChild(1));
		if (! descriptor.getType().isArray()) {
			throw new LDDCompilationException("" + tree.getText() 
					+ " not an array identifier " + createLineInformation(tree));
		}
		return new ArrayAccess(descriptor, indexExpression);
	}
	
	RVState getRVStateForName(RVClass rvClass, String stateName) {
		for (RVState rvState : rvClass.getDomain()) {
			if (rvState.getUid().getValue().equals(stateName)) {
				return rvState;
			}
		}
		
		return null;
	}
    
	Expression handleLiteral(Tree tree) {
		String textLiteral = tree.getText();
		Object literal = null;
		switch(tree.getType()) {
		case LDDParser.DecimalLiteral:
			// we are not sure if the number the number is small enought to use ints
			// instead of longs
			long longLiteral = Long.parseLong(textLiteral);
			if (longLiteral > Integer.MAX_VALUE || longLiteral < Integer.MIN_VALUE) { 
				// we ignore the fact that the value may be actually to big for long as well
				literal = longLiteral;				
			} else {
				literal = Integer.parseInt(textLiteral);
			}
			break;
			
		case LDDParser.FloatingPointLiteral:
			literal = Double.parseDouble(textLiteral);
			break;
		case LDDParser.TRUE:
			literal = Boolean.TRUE;
			break;
		case LDDParser.FALSE:
			literal = Boolean.FALSE;
			break;
		case LDDParser.ABSURD:
			// absurd is not expected as a literal expression in the
			// context required by this method
		default:
			throw createUnexpectedException(tree);
		}
		return new LiteralExpression(literal);
	}
	
	Block handleBlock(Tree tree) throws LDDCompilationException {
		List<Action> actions = new ArrayList<Action>();
		for (int i=0; i < tree.getChildCount(); i++) {
			actions.add((Action) handleAction(tree.getChild(i)));
		}
		return new Block(actions);
	}

	Action handleConditionalAction(Tree tree) throws LDDCompilationException {
		
		Expression condition = handleExpression(tree.getChild(0).getChild(0));
		Action ifAction = (Action) handleAction(tree.getChild(1));
		Action elseAction = null;
		if (tree.getChildCount() == 3 ) {
			elseAction = (Action) handleAction(tree.getChild(2));
		}
		if (condition.getValueType() != Boolean.class) {
			throw new LDDCompilationException(
					"Condition must be a boolean expression " + createLineInformation(tree.getChild(0)));
		}
		
		return new ConditionalAction(condition, ifAction, elseAction);
	}

	Descriptor handleIdentifier(Tree tree) throws LDDCompilationException {
    	String identifierName = null;
    	if (tree.getChildCount() == 1) {
    		// variables, random variables
    		identifierName = tree.getChild(0).getText();
    	} else 
    	if (tree.getChildCount() == 2) {
    		// random variable states
    		identifierName = tree.getChild(0).getText() + "." + tree.getChild(1).getText();
    		
    	} else {
    		throw new LDDCompilationException("Unrecognized identifier: "+ identifierName + " " + createLineInformation(tree));    	
    	}
    
		Descriptor descriptor = allDescriptorMap.get(identifierName);
		if (descriptor == null) {
			throw new LDDCompilationException("Undefined symbol " 
					+ identifierName + " " + createLineInformation(tree));
		}
		return descriptor;
	}
	

	FunctionCall handleFunctionCall(Tree tree) throws LDDCompilationException {
		Tree identifierChild = tree.getChild(0);
		assertTreeToken("function identifier", identifierChild, LDDParser.IDENTIFIER);
		Tree argumentChild = tree.getChild(1);
		assertTreeToken("function arguments", argumentChild, LDDParser.ARGS);
		
		Tree functionNameTree = identifierChild.getChild(0);
		String functionName = functionNameTree.getText();
		
		Class<?>[] parameterTypes = new Class<?>[argumentChild.getChildCount()];
		Expression[] expressions = new Expression[parameterTypes.length];
		for (int i=0; i < argumentChild.getChildCount(); i++) {
			Expression expr = (Expression) handleExpression(argumentChild.getChild(i));
			parameterTypes[i] = expr.getValueType();
			expressions[i] = expr;
		}
		Function function = FunctionRegistry.findClosestFunction(functionName, parameterTypes);		
		
		if (function == null) {
			throw new LDDCompilationException("Couln't find function for name: " 
					+ functionName + " and parameters: " + Arrays.asList(parameterTypes));
		}
		
		return new FunctionCall(function, expressions);
	}
	
	Expression handleBinaryOperation(Tree tree) throws LDDCompilationException {
		Expression expr1 =  (Expression) handleExpression(tree.getChild(0));
		Expression expr2 =  (Expression) handleExpression(tree.getChild(1));

		// TODO: add argument checks for binary expressions
		BinaryOperator op = null;
		switch(tree.getType()) {
		case LDDParser.AND:
			op = new LogicalAndOperator();
			break;
		case LDDParser.OR:
			op = new LogicalOrOperator();
			break;
		case LDDParser.PLUS: 
			op = new PlusOperator(expr1.getValueType(), expr2.getValueType());
			break;
		case LDDParser.MINUS: 
			op = new MinusOperator(expr1.getValueType(), expr2.getValueType());
			break;
		case LDDParser.TIMES: 
			op = new TimesOperator(expr1.getValueType(), expr2.getValueType());
			break;
		case LDDParser.DIV: 
			op = new DivOperator(expr1.getValueType(), expr2.getValueType());
			break;
		case LDDParser.REST: 
			op = new RestOperator(expr1.getValueType(), expr2.getValueType());
			break;
		case LDDParser.LT:	
			op = new LowerThanOperator(expr1.getValueType(), expr2.getValueType());
			break;
		case LDDParser.LTE:	
			op = new LowerEqualOperator(expr1.getValueType(), expr2.getValueType());
			break;
		case LDDParser.GT:	
			op = new GreaterThanOperator(expr1.getValueType(), expr2.getValueType());
			break;
		case LDDParser.GTE:	
			op = new GreaterEqualOperator(expr1.getValueType(), expr2.getValueType());
			break;
		case LDDParser.EQUALS:
			op = new EqualityOperator(expr1.getValueType(), expr2.getValueType());
			break;
			
		default: throw createUnexpectedException(tree);
		}
		return new BinaryOperation(expr1, expr2, op);
	}
	
	Class<?> checkTypesSame(Expression expr1, Expression expr2, Tree tree) throws LDDCompilationException {
		if (expr1.getValueType() != expr2.getValueType()) {
			throw new LDDCompilationException("Expressions have different types " + createLineInformation(tree));
		}
		return expr1.getValueType();
	}
	
	String[] getLeavesAsStrings(Tree tree) {
		String[] retString = new String[tree.getChildCount()];
		for(int i=0; i< retString.length; i++) {
			retString[i] = tree.getChild(i).getText();
		}
		return retString;
	}
	
	void handleParents(Tree parentsTree) throws LDDCompilationException {
		
		for (int i=0; i< parentsTree.getChildCount(); i++){
			Tree parentTree = parentsTree.getChild(i);
			String parentName = parentTree.getText();
			if (rvClassesByNameMap.containsKey(parentName)) {
				if (this.parentsByNameMap.containsKey(parentName)) {
					throw new LDDCompilationException(
							"Duplicate parent " + parentName 
							+ " " + createLineInformation(parentTree));
				}
				RVClass parentClass = rvClassesByNameMap.get(parentName);
				this.parentsByNameMap.put(parentName, parentClass);
				this.parents.add(parentClass);

				RVClass rvClass = rvClassesByNameMap.get(parentName);
				addDescriptor(parentTree, new DescriptorImpl(
						parentName, 
						rvClass.getClass(),
						DescriptorType.RandomVariable, 
						false, rvClass
						), false);
				for(RVState state: rvClass.getDomain()) {
					addDescriptor(parentTree, new DescriptorImpl(
							state.getRVClass().getName() + "." + state.getUid().getValue(), 
							state.getClass(),
							DescriptorType.RandomVariableState, 
							false, state
							), false);
					
				}
				
			} else {
				throw new LDDCompilationException(
						"Parent " + parentName 
						+ " unknown " + createLineInformation(parentTree));
			}
		}
		
	
	}
	
	void handleDistributionId(Tree tree) {
		if (tree.getChildCount() != 0) {
			distributionId = tree.getChild(0).getText();
		}
	}
	
	String createLineInformation(Tree tree) {
		return "in line: " + tree.getLine() + ", col: " + tree.getCharPositionInLine() + ".";
	}
	
	void handleDomain(Tree domainTree) throws LDDCompilationException {
		for (int i=0; i< domainTree.getChildCount(); i++) {
			Tree uidTree = domainTree.getChild(i);
			String uidName ;
			if (uidTree.getType() == LDDParser.ABSURD) {
				uidName = MEBNConstants.ABSURD;
			} else if (uidTree.getType() == LDDParser.TRUE) {
				uidName = MEBNConstants.TRUE;
			} else if (uidTree.getType() == LDDParser.FALSE) {
				uidName = MEBNConstants.FALSE;
			} else {
				uidName = uidTree.getText();
			}

			if (this.allUidsByNameMap.containsKey(uidName)) {
				if (this.domainUidsByName.containsKey(uidName)) {
					throw new LDDCompilationException(
							"Duplicate UID " + uidName + " in domain definition "
							+ createLineInformation(uidTree));
				}
				this.domainUidsByName.put(uidName, this.allUidsByNameMap.get(uidName));
				this.domain.add(this.allUidsByNameMap.get(uidName));
			} else {
				throw new LDDCompilationException(
						"UID " + uidName 
						+ " unknown in domain definition " 
						+ createLineInformation(uidTree));
			}
		}
		if (! this.domainUidsByName.containsKey(MEBNConstants.ABSURD)) {
			throw new LDDCompilationException("Absurd value expected in domain specification " 
					+ createLineInformation(domainTree));
		}
		
		
		// create the distribution array
		addDescriptor(domainTree, 
				new DescriptorImpl(RETURN_ARRAY_NAME, Double[].class, DescriptorType.Variable, true),
				true);
	}
	
	private IllegalArgumentException createUnexpectedException(Tree tree) {
		return new IllegalArgumentException("Unexpected token: " + tree.getText() + " (" + tree.getType() + ").");
	}
	
	
	void assertTreeToken(String description, Tree tree, int token) {
		if (tree.getType() != token) {
			throw createUnexpectedException(tree);
		}
	}
	
	
}